# 线性回归
"""
前馈计算：从前往后计算loss
反向传播：从后向前利用链式求导法则计算梯度，然后使用得到的梯度更新参数

pytorch 中用于计算的变量通常是一个张量 tensor，tensor.data 才是具体的值，tensor.grad 是该参数的梯度，pytorch 中 tensor 之间的计算会自动构建计算图，
计算图中的 tensor 可以使用 backward() 自动求梯度。

pytorch 中构建神经网络模型都要继承自 nn.Module 模块，该模块中包含了构建神经网络的一些结构。

构建神经网路训练流程：
    1、准备数据集
    2、定义模型类，继承nn.Module，
        2.1、实现 __init__ 函数，用于变量的初始化，第一行调用父类的 __init__
        2.2、定义前馈计算 forward() 函数
    3、构造损失函数和优化器 optimizer
    4、训练
        4.1、定义训练轮数 Epochs
            4.1.1、计算预测值 y'
            4.1.2、计算损失loss，这里会构建计算图
            4.1.3、方向传播 backward() 函数求梯度，这里会释放计算图，不会每次都生成一个计算图去占用资源
            4.1.4、更新参数，这里使用 w.grad.data
            4.1.5、梯度清0（可选操作），防止梯度累加
        4.2、结束循环
    5、训练结束
"""
import torch
import torch.nn as nn
import matplotlib.pyplot as plt

# 给定假设数据,，4个一维样本向量
x_data = torch.Tensor([[1.0], [2.0], [3.0], [4.0]])
y_data = torch.Tensor([[2.0], [4.0], [6.0], [8.0]])


# 线性回归模型
class LinearModel(nn.Module):
    def __init__(self):
        """
        这里面通常会包括一些需要训练的参数，在这里进行声明并初始化
        """
        # 调用符类的 __init__
        super(LinearModel, self).__init__()
        # Linear(1, 1) 有三个主要参数，in_features: int，输入特征维度；out_features: int，输出特征维度 bias: bool = True，是否使用偏执项
        self.linear = torch.nn.Linear(1, 1)

    def forward(self, x):
        """
        前馈计算 y'

        Args:
            x:数据集中的输入
        """
        y_pred = self.linear(x)
        return y_pred


# 获得模型
model = LinearModel()

# 构建损失函数，使用平均平方误差 MSE，size_average指是否使用loss的均值
criterion = torch.nn.MSELoss(size_average=False)
# 构建优化器，使用随机梯度下降，model.parameters()是将模型需要更新的参数传给优化器进行更新, lr=0.01 是学习率
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 迭代训练
for epoch in range(100):
    # 前馈计算，多个样本同时计算，通常是使用矩阵进行计算
    y_pred = model(x_data)
    # 计算损失
    loss = criterion(y_pred, y_data)
    # 输出迭代轮次和损失
    print(epoch, loss)
    # 梯度清 0
    optimizer.zero_grad()
    # 方向计算梯度
    loss.backward()
    # 使用优化器的方法更新参数
    optimizer.step()

# 输出权重和偏置项
print('w = ', model.linear.weight.item())
print('b = ', model.linear.bias.item())

# 测试集
x_test = torch.Tensor([[4.0]])
# 得到预测值
y_test = model(x_test)
# 输出预测值
print('y_pred = ', y_test.data)
